from model.Densenet import Densenet
from model.MobileNet import MobileNet
from util import Util, F1_score
from tensorflow.python.keras.metrics import Precision, Recall, AUC
import pandas as pd
from tensorflow.python.keras.callbacks import EarlyStopping
import tensorflow as tf


class Train(Util):

    def model_select(self):
        if self.model_name.__contains__('DenseNet'): return Densenet().model
        elif self.model_name.__contains__('MobileNet'): return MobileNet().model

    def model_train(self):
        train_ds, val_ds = self.get_data()
        model = self.model_select()
        model.compile(optimizer='adam',
                      loss='categorical_crossentropy',
                      metrics=['accuracy', F1_score(), Recall(), Precision(), AUC()])
        early_stopping = EarlyStopping(
            monitor=self.monitor,
            verbose=1,
            patience=self.patience,
            restore_best_weights=True
        )
        reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(min_lr=0.00001,
                                                         factor=0.2)

        history = model.fit(train_ds, epochs=2000, callbacks=[early_stopping, reduce_lr], validation_data=val_ds)

        hist_df = pd.DataFrame(history.history)
        hist_df.to_csv(self.model_name + '.csv')
        model.save(self.model_name + '.h5')

if __name__ == '__main__':
    train = Train()
    train.model_train()
